package odesimulation;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;

import javax.swing.JFrame;

import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.misc.TestRig;

import odesimulation.testantlrParser.ReactionsetContext;
import optimizers.commons.GAConstants;

import org.apache.commons.math3.ode.FirstOrderDifferentialEquations;
import org.apache.commons.math3.ode.nonstiff.ClassicalRungeKuttaIntegrator;


import util.MyStepHandler;
import visualization.PlotExpData;

public class OdeSolver implements FirstOrderDifferentialEquations {

	public final static ArrayList<String> names= new ArrayList<String>(Arrays.asList(new String[]{"x","y","z"})); // x, y, z...
	public String[] myEquations; // expressed in polish notation
	public HashMap<String,Double> params; // contains the value of a, b, c...
	public final static ArrayList<String> possibleOps = new ArrayList<String>(Arrays.asList(new String[]{"+","-","*","/"}));
	
	public static int numberOfSteps = 100;
	
	public OdeSolver(String[] myEquations, HashMap<String,Double> params){
		this.myEquations = myEquations;
		this.params = params;
	}
	
	@Override
	public void computeDerivatives(double t, double[] y, double[] ydot) {
		for(int i = 0; i < myEquations.length; i++){
			ydot[i] = parse(y,myEquations[i]);
			System.out.println(names.get(i)+": "+ydot[i]);
		}

	}

	private double parse(double[] y, String string) {
		char elem;
		ArrayList<String> operationList = new ArrayList<String>();
		for(int i=0; i< string.length(); i++){
			elem = string.charAt(i);
			
			if(elem == '(' || elem == ')' ||  elem == ' '){
				continue; // polish notation is unambiguous with operators of known arity
			}
			if(possibleOps.contains(""+elem)){
				operationList.add(""+elem);
			} else {
				//hopefully a known character
				if(names.contains(""+elem) || params.containsKey(""+elem)){
				    String prev = "";
					if (operationList.size()>0)
						prev = operationList.get(operationList.size()-1);
					
					double myvalue = names.contains(""+elem)?y[names.indexOf(""+elem)]:params.get(""+elem);
					while(operationList.size()>1 && !possibleOps.contains(prev)){
					//we got a complete operation!
					double prevvalue = Double.parseDouble(operationList.get(operationList.size()-1));
					switch(operationList.get(operationList.size()-2).charAt(0)){
					case '+':
						prevvalue += myvalue;
						break;
					case '-':
						prevvalue -= myvalue;
						break;
					case '*':
						prevvalue *= myvalue;
						break;
					case '/':
						prevvalue /= myvalue;
						break;
					default:
						System.err.println("Can't parse equation");
						return 0;
					}
					
					operationList.remove(operationList.size()-1);
					operationList.remove(operationList.size()-1);
					myvalue = prevvalue;
					if(operationList.size() > 0){
						prev = operationList.get(operationList.size()-1);
					}
					
				}
					operationList.add(""+myvalue);
					continue;
				}
				System.err.println("Can't parse equation");
				return 0;
			}
		}
		return Double.parseDouble(operationList.get(0));
	}
	
	@Override
	public int getDimension() {
		return myEquations.length;
	}

	public static double[][] solveEqus(double[] initialConditions, String[] myEquations, HashMap<String,Double> params){
		OdeSolver odes = new OdeSolver(myEquations, params);
		ClassicalRungeKuttaIntegrator myIntegrator = new ClassicalRungeKuttaIntegrator(1e-1); 
		MyStepHandler handler = new MyStepHandler(odes.getDimension(), null);
		myIntegrator.addStepHandler(handler);
		myIntegrator.integrate(odes, 0, initialConditions, OdeSolver.numberOfSteps,
				initialConditions);
		return handler.getTimeSerie();
	}
	
	//public static A
	
	public static void main(String[] args){
		testantlrLexer myLex = new testantlrLexer(new ANTLRInputStream("a -> b;"));
		testantlrParser myPars = new testantlrParser(new CommonTokenStream(myLex));
		ReactionsetContext context = myPars.reactionset();
		System.out.println(context.toStringTree()+context.symbols.get(0));
		HashMap<String,Double> params = new HashMap<String,Double>();
		params.put("a", 1.5);
		params.put("b", 3.0);
		params.put("c", 3.0);
		params.put("d", 5.0);
		params.put("0", 0.0);
		String[] myEqus = new String[]{"- * * a x x * * b x y", "- * * b x y * c y", "0"};
		OdeSolver odes = new OdeSolver(myEqus,params);
		double[] values = new double[]{1.0,0.1,10.1};
		System.out.println(odes.parse(values, myEqus[0]));
		model.Constants.numberOfPoints = OdeSolver.numberOfSteps;
		double[][] data = OdeSolver.solveEqus(values, myEqus, params);
		String[] seqNames = new String[OdeSolver.names.size()];
		seqNames = OdeSolver.names.toArray(seqNames);
		PlotExpData demo = new PlotExpData("Solved equation",
				data, seqNames, "Time", "Expression");
		JFrame frame = new JFrame();
		frame.add(demo);
		frame.pack();
		frame.setVisible(true);
		frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
	}
	
}
